package com.lktx.sso.scope;


import cn.dev33.satoken.context.SaHolder;
import cn.dev33.satoken.context.model.SaRequest;
import cn.dev33.satoken.oauth2.SaOAuth2Manager;
import cn.dev33.satoken.oauth2.data.model.oidc.IdTokenModel;
import cn.dev33.satoken.oauth2.data.model.request.ClientIdAndSecretModel;
import cn.dev33.satoken.oauth2.scope.handler.OidcScopeHandler;
import cn.dev33.satoken.stp.StpUtil;
import cn.hserver.core.ioc.annotation.Autowired;
import cn.hserver.core.ioc.annotation.Bean;
import cn.hserver.core.ioc.annotation.Value;
import cn.hserver.plugin.web.context.HServerContextHolder;
import cn.hserver.plugin.web.context.HttpSession;
import cn.hutool.core.util.StrUtil;
import cn.hutool.crypto.asymmetric.RSA;
import cn.hutool.json.JSONUtil;
import cn.hutool.jwt.JWT;
import cn.hutool.jwt.signers.JWTSigner;
import cn.hutool.jwt.signers.JWTSignerUtil;
import com.lktx.sso.admin.dto.KvDTO;
import com.lktx.sso.admin.entity.SsoAppOidcConfig;
import com.lktx.sso.admin.mapper.SsoUserMapper;
import com.lktx.sso.admin.service.SsoAppOidcConfigService;
import com.lktx.sso.config.OAuthConfig;

import java.util.HashMap;
import java.util.Map;


@Bean
public class CustomOidcScopeHandler extends OidcScopeHandler {

    @Autowired
    private SsoAppOidcConfigService ssoAppOidcConfigService;

    @Autowired
    private SsoUserMapper ssoUserMapper;

    @Value("oauth.baseIss")
    private String baseIss;

    @Override
    public String getIss() {
        SaRequest req = SaHolder.getRequest();
        ClientIdAndSecretModel client = SaOAuth2Manager.getDataResolver().readClientIdAndSecret(req);
        SsoAppOidcConfig byClientId = ssoAppOidcConfigService.getByClientId(client.clientId);

        return String.format("%s/%s", baseIss, byClientId.getSsoAppId());
    }

    @Override
    public IdTokenModel workExtraData(IdTokenModel idToken) {
        Map<String, Object> userInfo = ssoUserMapper.getLoginUserInfo(idToken.sub.toString());

        System.out.println("workExtraData:" + JSONUtil.toJsonStr(userInfo));
        System.out.println("----- 为 idToken 追加扩展字段 ----- ");
        idToken.extraData.putAll(userInfo);
        // 更多字段 ...
        // 可参考：https://openid.net/specs/openid-connect-core-1_0.html#StandardClaims
        return idToken;
    }

    @Override
    public String generateJwtIdToken(IdTokenModel idToken) {
        SsoAppOidcConfig byClientId = ssoAppOidcConfigService.getByClientId(idToken.azp);
        if (byClientId == null) {
            throw new RuntimeException("OIDC 配置错误");
        }
        Map<String, Object> dataMap = convertIdTokenToMap(idToken);
//        String code = SaHolder.getRequest().getParam("code");
//        String nonce = OAuthConfig.CACHE_CODE_NONCE.getIfPresent(code);
//        if (StrUtil.isNotEmpty(nonce)) {
//            dataMap.put("nonce", nonce);
//        }
        JWT jwt = JWT.create().addPayloads(dataMap);
        if (byClientId.getJwtPayload() != null) {
            for (KvDTO kvDTO : byClientId.getJwtPayload()) {
                jwt.getPayload().setPayload(kvDTO.getKey(), kvDTO.getValue());
            }
        }
        if (byClientId.getJwtHeader() != null) {
            Map<String, Object> data = new HashMap<>();
            for (KvDTO kvDTO : byClientId.getJwtHeader()) {
                data.put(kvDTO.getKey(), kvDTO.getValue());
            }
            jwt.getHeader().addHeaders(data);
        }
        jwt.getHeader().setKeyId(idToken.azp);
        RSA rsa = new RSA(byClientId.getPrivateKey(), null);
        JWTSigner jwtSigner;
        if (byClientId.getIdTokenSignatureAlgorithm().equals("RS256")) {
            jwtSigner = JWTSignerUtil.rs256(rsa.getPrivateKey());
        } else if (byClientId.getIdTokenSignatureAlgorithm().equals("ES256")) {
            jwtSigner = JWTSignerUtil.es256(rsa.getPrivateKey());
        } else {
            jwtSigner = JWTSignerUtil.hs256(byClientId.getPublicKey().getBytes());
        }
        return jwt.setSigner(jwtSigner).sign();
    }
}